"""
Author: Morphlng
Date: 2024-03-23 13:29:25
LastEditTime: 2024-03-25 14:16:03
LastEditors: Morphlng
Description: Test for metadrive adapter
FilePath: /DrivingGym/tests/metadrive/metadrive_adapter_test.py
"""

import numpy as np
from metadrive.component.sensors.rgb_camera import RGBCamera
from metadrive.component.sensors.semantic_camera import SemanticCamera

from driving_gym.misc.util import norm_angle
from driving_gym.simulation.adapter_factory import AdapterFactory

env_cfg = dict(
    num_agents=5,
    log_level=50,
    image_observation=True,
    sensors=dict(
        sementic_camera=[SemanticCamera, 200, 100], rgb_camera=[RGBCamera, 200, 100]
    ),
    vehicle_config={"image_source": "sementic_camera"},
)
adapter = AdapterFactory.create(env_cfg, adapter_type="metadrive")


class TestAdapter:
    def test_env_initialize(self):
        adapter.initialize()
        assert adapter.env is not None, "Environment is not initialized"

    def test_load_scenario(self):
        assert adapter._last_obs is None, "Adapter isn't cleaned up"
        adapter.load_scenario(adapter.config)
        assert adapter._last_obs is not None, "Scenario is not loaded"
        assert (
            len(adapter.agents) == env_cfg["num_agents"]
        ), "Agents should be available after loading scenario"

    def test_get_transform(self):
        trans = adapter.get_transform("agent0")
        engine_trans = adapter.env.agents["agent0"].origin.get_transform()
        np.testing.assert_allclose(
            (trans.location.x, trans.location.y, trans.location.z),
            (engine_trans.pos.x, engine_trans.pos.y, engine_trans.pos.z),
            err_msg="Transform location is not correct",
        )

        np.testing.assert_allclose(
            (trans.rotation.yaw, trans.rotation.pitch, trans.rotation.roll),
            (
                norm_angle(engine_trans.hpr.x + 90),
                norm_angle(engine_trans.hpr.y),
                norm_angle(engine_trans.hpr.z),
            ),
            err_msg="Transform rotation is not correct",
        )

    def test_get_velocity(self):
        velocity = adapter.get_velocity("agent0")
        engine_velocity = adapter.env.agents["agent0"].velocity
        np.testing.assert_allclose(
            (velocity.x, velocity.y),
            engine_velocity,
            err_msg="Velocity is not correct",
        )

    def test_get_acceleration(self):
        env = adapter.env

        for _ in range(100):
            actions = env.action_space.sample()

            for agent in env.agents:
                acceleration = adapter.get_acceleration(agent)
                assert np.isclose(
                    acceleration.x + 1, adapter._last_info[agent]["acceleration"] + 1
                ), "Acceleration is not correct"
                adapter.apply_control(agent, actions[agent])

            adapter.tick()

    def test_bounding_box(self):
        bbox = adapter.get_bounding_box("agent0")
        trans = adapter.get_transform("agent0")
        # Metadrive ignores z-axis
        trans.location.z = 0

        global_bbox = bbox.transform_to_world_frame(trans, True)[:, :2]
        engine_bbox = np.array(adapter.env.agents["agent0"].bounding_box)

        # TODO: if roll and pitch are both not 0, assert will fail
        np.testing.assert_allclose(
            global_bbox,
            np.roll(engine_bbox[::-1, :], 1, axis=0),
            rtol=1e-3,
            err_msg="Bounding Box does not match",
        )

    def test_sensor_data(self):
        env = adapter.env
        agents = env.agents
        sensors = env.engine.sensors

        for agent in agents:
            for sensor in sensors:
                data = adapter.get_sensor_data(f"{agent}:{sensor}")
                assert data is not None, "Sensor data is not available"

        # top_down view is not an actual sensor
        data = adapter.get_sensor_data(f"{agent}:top_down")
        assert data is not None, "Top down view is not available"

    def test_apply_control(self):
        env = adapter.env
        last_info = adapter._last_info.copy()
        actions = env.action_space.sample()
        agent_nums = len(env.agents)

        for i, agent in enumerate(env.agents):
            adapter.apply_control(agent, actions[agent])
            adapter.tick()

            if i < agent_nums - 1:
                assert (
                    adapter._last_info[agent].get("episode_length", 0)  # fmt: skip
                    == last_info[agent].get("episode_length", 0)
                ), "Episode length is not updated"
            else:
                assert (
                    adapter._last_info[agent].get("episode_length", 0)  # fmt: skip
                    == last_info[agent].get("episode_length", 0) + 1
                ), "Episode length should be updated"

    def test_cleanup(self):
        adapter.cleanup()
        assert adapter.env is None, "Environment is not cleaned up"
